{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Simulation\n",
    "\n",
    "In this tutorial, we will show how to use prysm to perform very accurate image simulation.  This combines previously shown concepts in the [Lens MTF Model](./Lens-MTF-Model.ipynb) with a few additional ones and is a more advanced tutorial.  The reader is assumed to understand basic propagation for computing PSFs.\n",
    "\n",
    "Any image chain model or image simulation begins with defining parameters, so we'll choose some fairly arbitrary ones for our system:\n",
    "\n",
    "Detector:\n",
    "- pixel pitch: 4.5 microns\n",
    "- Output resolution: 512x512\n",
    "- Read Noise: 10 e-\n",
    "- Full-well capacity: 50,000 e-\n",
    "- Dark current: 30 e-/px/s\n",
    "- bias: 800 e-\n",
    "- conversion gain 5e-/DN\n",
    "- Bit depth: 12-bit\n",
    "\n",
    "Optics:\n",
    "- lens F/#: 2.8\n",
    "- lens EFL: 100 mm\n",
    "- aperture: circular\n",
    "- Optical path error? yes - sum of 2D-Q polynomials\n",
    "- Fully achromatic (constant OPD over all wavelengths)\n",
    "\n",
    "Object/Scene:\n",
    "- Object: Siemens' Star\n",
    "- Spectrum: Gaussian about 550 nm, 10% fractional bandwidth\n",
    "\n",
    "From these, we begin to determine the parameters of the forward model.  The model _must_ be at least Nyquist sampled, or due to aliasing it will be invalid.  We define:\n",
    "\n",
    "$$\n",
    "Q = \\frac{\\lambda \\text{F\\#}}{pp}\n",
    "$$\n",
    "where $pp$ is the pixel pitch, and compute $Q = \\tfrac{2.8 * .495}{4.5} = 0.306$.  495 nm is 10% below 550 nm.  Since we require $Q>=2$ in the forward model, and assume at the moment that we will use an integer level of oversampling, then the forward model is run at $Q=\\text{roundup}\\{2/.306\\}=7x$ oversampling, or $Q=2.156$.  A notable problem is that this equation for $Q$ contains the wavelength; in other words, $Q$ is chromatic.  To get around this, we'll use matrix triple product DFTs to propagate directly to the oversampled version of the detector grid, which will be 3584x3584 samples across."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy import stats\n",
    "\n",
    "from prysm import (\n",
    "    coordinates,\n",
    "    convolution,\n",
    "    detector,\n",
    "    geometry,\n",
    "    propagation,\n",
    "    polynomials,\n",
    "    objects,\n",
    ")\n",
    "\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we define a bunch of parameters and set up the basic representation of the pupil:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp = 4.5\n",
    "res = 512\n",
    "fno = 2.8\n",
    "efl = 100\n",
    "epd = efl/fno\n",
    "r_aper = epd / 2\n",
    "\n",
    "xi, eta = coordinates.make_xy_grid(1800, diameter=epd)\n",
    "r, t = coordinates.cart_to_polar(xi,eta)\n",
    "dx = xi[0,1] - xi[0,0]\n",
    "\n",
    "r_aber = r / r_aper\n",
    "\n",
    "amp = geometry.circle(r_aper, r)\n",
    "\n",
    "nms = [polynomials.noll_to_nm(j) for j in range(1,11)]\n",
    "basis = list(polynomials.Q2d_sequence(nms, r_aber, t))\n",
    "\n",
    "phs_coefs = np.random.rand(len(basis)) * 2000 # 200/sqrt(12) nm per mode, per uniform distribution statistics\n",
    "\n",
    "phs = polynomials.sum_of_2d_modes(basis, phs_coefs)\n",
    "\n",
    "# only used for plotting\n",
    "mask = amp == 0\n",
    "phs2 = phs.copy()\n",
    "phs2[mask] = np.nan\n",
    "im = plt.imshow(phs2, cmap='inferno')\n",
    "plt.colorbar(im)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note carefully the use of r_aber when computing the basis -- the polynomials are defined on $r \\in [0,1]$, which necessitates this normalization.  We've used the smallest possible grid that can contain the pupil with no padding, and a large number of samples.  A requirement of the matrix triple product DFT is that the output resolution, divided by Q, must not exceed the input resolution.  If this is not honored, then Dirichlet clones of the PSF will be visible in the edge of the array, which is unphysical and incorrect.  Since we will have 3584 px on the output and Q~=2.1, we need about 1800 px on the input,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wvl0 = .550\n",
    "halfbw = 0.1\n",
    "wvls = np.linspace(wvl0*(1-halfbw), wvl0*(1+halfbw), 7)\n",
    "\n",
    "def gauss(x, mu, sigma):\n",
    "    num = (x-mu)**2\n",
    "    den = 2 * sigma ** 2\n",
    "    return np.exp(-num/den)\n",
    "\n",
    "spectral_weights = gauss(wvls, .550, .550*.05)\n",
    "plt.stem(wvls, spectral_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having now completed the bulk of the preparatory work, we can now compute the PSF associated with this OPD and this spectrum.  We'll propagate each wavelength to the oversampled grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "psfs = []\n",
    "for wvl in wvls:\n",
    "    pup = propagation.Wavefront.from_amp_and_phase(amp, phs, wvl, dx)\n",
    "    tmp = pup.focus_fixed_sampling(efl, pp/7, 3584)\n",
    "    psfs.append(tmp.intensity.data)\n",
    "\n",
    "# re-use this function, identical behavior\n",
    "psf = polynomials.sum_of_2d_modes(psfs, spectral_weights)\n",
    "\n",
    "# norm to sum of 1, no loss or creation of energy from convolution\n",
    "psf /= psf.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After producing the PSF, we need to rasterize the target.  For this, we need a new grid:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = coordinates.make_xy_grid(tmp.data.shape, dx=tmp.dx)\n",
    "r, t = coordinates.cart_to_polar(x,y)\n",
    "obj = objects.siemensstar(r,t, 100, oradius=x.max()*.8)\n",
    "\n",
    "plt.figure(figsize=(15,15))\n",
    "plt.imshow(obj, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now in posession of the object and PSF, we can synthesize the aerial image, which has been blurred by the optical system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = convolution.conv(obj, psf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The blur by the detector is modeled in this case as a 100% fill factor pixel.  Binning will bring us to the output reoslution and simultaneously:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img2 = detector.bindown(img, 7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To treat noise, we build a detector model, and capture the image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peak_eminus = 45_000 # heavy underexposure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dark_current = 10\n",
    "fwc = 50_000\n",
    "rn = 10\n",
    "bias = 800\n",
    "kgain = 5\n",
    "bit_depth = 12\n",
    "texp = 1/60 # 1/60 sec\n",
    "cam = detector.Detector(dark_current, rn, bias, fwc, kgain, bit_depth, texp)\n",
    "\n",
    "im = cam.expose(img2*peak_eminus)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15,15))\n",
    "plt.imshow(im, cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this, we have fully modeled the image chain, including optical and electrical blurs and noise.  "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}